# Send message to log
log_code()

stan_data <- list(
  C = length(c_id),
  P = length(p_id),
  I = length(i_id),
  E = length(e_id),
  F1 = 1,
  F2 = 1,
  p = df_long$value,
  N = nrow(df_long),
  e_id = df_long$e_id,
  c_id = df_long$c_id,
  p_id = df_long$p_id,
  i_id = df_long$i_id,
  free_experts = 0,
  country_issue_shocks = 0,
  party_issue_shocks = 0,
  country_shocks = 0,
  f1 = c(1:dim(df[,7:27])[2])[!(1:dim(df[,7:27])[2] %in% c(1))],
  f2 = c(1:dim(df[,7:27])[2])[!(1:dim(df[,7:27])[2] %in% c(7))],
  f1z = array(c(1), c(1,1)), # constrained issue
  f2z = array(c(7), c(1,1)) # constrained issue
)

# setting up starting values
initF <- function(){
  list(theta_free = as.matrix(p_start[,c("lrecon","galtan")], ncol = 2),
       mu_alpha = matrix(seq(-2, 2, length.out = 10), stan_data$I, 10, byrow = TRUE),
       hyp_p = c(1, 1, .1, .1, .1, .1),
       res_beta_r1 = matrix(0, stan_data$I-stan_data$F1, stan_data$C),
       res_beta_r2 = matrix(0, stan_data$I-stan_data$F2, stan_data$C),
       res_alpha_r = matrix(0, stan_data$I, stan_data$C),
       scales_beta_1 = rep(1/(stan_data$I-stan_data$F1), stan_data$I-stan_data$F1),
       scales_beta_2 = rep(1/(stan_data$I-stan_data$F2), stan_data$I-stan_data$F2),
       scales_alpha = rep(1/stan_data$I, stan_data$I),
       L_beta_1 = chol(diag(rep(1, stan_data$I-stan_data$F1))),
       L_beta_2 = chol(diag(rep(1, stan_data$I-stan_data$F2))),
       L_alpha = chol(diag(rep(1, stan_data$I))),
       exp_scale_free = brms::rdirichlet(1, rep(100, stan_data$E)),
       scale_beta_free = brms::rdirichlet(1, rep(100, stan_data$C)),
       shift_alpha_free = (rnorm(stan_data$C, 0, .001)),
       idio_r = matrix(0, nrow = stan_data$I, ncol = stan_data$P))
}

mod <- cmdstan_model("full_model.stan")

# baseline model ####
fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_m0 <- loo(fit$draws("log_lik"))
pred_m0 <- fit$draws("p_share")
df_m0 <- fit$summary(c("mu_beta_1", "mu_beta_2", "mu_alpha", "theta", "p_share"))
save(df_m0, loo_m0, pred_m0, file = "out/m0.Rda")

gc();rm(fit);Sys.sleep(30);gc()

# heteroskedastic experts model ####
stan_data$free_experts <- 1

fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_m1 <- loo(fit$draws("log_lik"))
pred_m1 <- fit$draws("p_share")
df_m1 <- fit$summary(c("mu_beta_1", "mu_beta_2", "mu_alpha", "theta", "p_share", "exp_scale"))
save(df_m1, loo_m1, pred_m1, file = "out/m1.Rda")

gc();rm(fit);Sys.sleep(30);gc()

# idiosyncratic preferences model ####
stan_data$free_experts <- 0
stan_data$party_issue_shocks <- 1

fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_m2 <- loo(fit$draws("log_lik"))
pred_m2 <- fit$draws("p_share")
df_m2 <- fit$summary(c("mu_beta_1", "mu_beta_2", "mu_alpha", "theta", "p_share", "LL_alpha"))
save(df_m2, loo_m2, pred_m2, file = "out/m2.Rda")

gc();rm(fit);Sys.sleep(30);gc()

# A-M model ####
stan_data$party_issue_shocks <- 0
stan_data$country_shocks <- 1

fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_m3 <- loo(fit$draws("log_lik"))
pred_m3 <- fit$draws("p_share")
df_m3 <- fit$summary(c("mu_beta_1", "mu_beta_2", "mu_alpha", "theta", "p_share", "LL_alpha", "shift_alpha", "scale_beta"))
save(df_m3, loo_m3, pred_m3, file = "out/m3.Rda")

gc();rm(fit);Sys.sleep(30);gc()

# country-issue shocks model ####
stan_data$country_shocks <- 0
stan_data$country_issue_shocks <- 1

fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_m4 <- loo(fit$draws("log_lik"))
pred_m4 <- fit$draws("p_share")
df_m4 <- fit$summary(c("beta_1", "beta_2", "mu_beta_1", "mu_beta_2", "mu_alpha", "theta", "p_share", "LL_alpha", "res_beta_1", "res_beta_2", "res_alpha"))
save(df_m4, loo_m4, pred_m4, file = "out/m4.Rda")

gc();rm(fit);Sys.sleep(30);gc()

# full model ####
stan_data$free_experts <- 1
stan_data$country_issue_shocks <- 1
stan_data$party_issue_shocks <- 1
stan_data$country_shocks <- 1

fit <- mod$sample(data = stan_data, 
                  init = initF,
                  chains = 8,
                  iter_warmup = 500, iter_sampling = 500)

loo_full <- loo(fit$draws("log_lik"))
pred_full <- fit$draws("p_share")
df_full <- fit$summary(c("beta_1", "beta_2", "mu_alpha", "mu_beta_1", "mu_beta_2", "theta", "res_beta_1", "res_beta_2", "res_alpha", "exp_scale", "LL_beta1", "LL_beta2", "LL_alpha", "idio", "shift_alpha", "scale_beta", "hyp_p", "p_share"))
save(loo_full, pred_full, df_full, file = "out/full1.Rda")

rm(list = setdiff(ls(), "lf"))